


from jittor.dataset import Dataset
import os
import cv2
import numpy as np

# IMG_SIZE=(224,224)
# IMG_SIZE=(299,299)
IMG_SIZE=(448,448)

import cv2


def resize_keep_aspectratio(image_src, dst_size):
    src_h, src_w = image_src.shape[:2]
    # print(src_h, src_w)
    dst_h, dst_w = dst_size

    # 判断应该按哪个边做等比缩放
    h = dst_w * (float(src_h) / src_w)  # 按照ｗ做等比缩放
    w = dst_h * (float(src_w) / src_h)  # 按照h做等比缩放

    h = int(h)
    w = int(w)

    if h <= dst_h:
        image_dst = cv2.resize(image_src, (dst_w, int(h)))
    else:
        image_dst = cv2.resize(image_src, (int(w), dst_h))

    h_, w_ = image_dst.shape[:2]
    # print(h_, w_)

    top = int((dst_h - h_) / 2);
    down = int((dst_h - h_ + 1) / 2);
    left = int((dst_w - w_) / 2);
    right = int((dst_w - w_ + 1) / 2);

    value = [0, 0, 0]
    borderType = cv2.BORDER_CONSTANT
    # print(top, down, left, right)
    image_dst = cv2.copyMakeBorder(image_dst, top, down, left, right, borderType, None, value)

    return image_dst

class Dogs(Dataset):
    def __init__(self, data_root="/home/hfle/data_dog/TEST_A", train=True , batch_size=1, shuffle=False):
        # if you want to test resnet etc you should set input_channel = 3, because the net set 3 as the input dimensions
        super().__init__()
        self.data_root = data_root
        self.batch_size = batch_size
        
        self.imgs=list()
        self.labels=list()

        file_names=os.listdir(data_root)
        for file_name in file_names:
            file_path=os.path.join(data_root,file_name)
            img=cv2.imread(file_path)
            img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
            img=resize_keep_aspectratio(img,IMG_SIZE)
            img=img.transpose((2,0,1))
            img=img.astype(np.float32)/255.0
            self.imgs.append(img)
            self.labels.append(file_name)

        self.total_len=len(self.imgs)
        # this function must be called
        self.set_attrs(batch_size = self.batch_size, total_len=self.total_len, shuffle= self.shuffle)


    def __getitem__(self, index):
        return self.imgs[index], self.labels[index]